from dataclasses import dataclass
from typing import Any
from json import JSONEncoder

import requests
import json
from itertools import combinations
import copy
import time

from zeus.operation_service.app.settings import configuration
from vulcanus.log.log import LOGGER
from vulcanus.restful.response import BaseResponse
from vulcanus.restful.resp.state import (
    SUCCEED,
    UNKNOWN_ERROR,
)

TEN_MINUTE = 600000
TIME_DEVIATION = 3000000000


@dataclass
class Span:
    name: str
    start_timestamp: int
    end_timestamp: int
    children: list
    tgid: int
    thread_id: int
    id: str
    client_ip: str
    server_ip: str
    client_port: int
    server_port: int
    protocol: str
    value: int
    req_type: str
    req_status: str
    msg: Any
    raw_data: Any

    @dataclass
    class HttpMsg:
        req_method: str
        req_path: str
        resp_status: int
        resp_message: str

    @dataclass
    class RedisMsg:
        req_command: str
        req_payload: str
        resp_command: str
        resp_payload: str

    @classmethod
    def parse_msg(cls, keys, s):
        values = s.split('@@')
        result = {}
        for i in range(min(len(keys), len(values))):
            result[keys[i]] = values[i]
        return result

    @classmethod
    def parse_span(cls, data):
        msg = None
        name = 'unknown'
        req_type = '-'
        req_status = '-'
        if data['protocol'] == 'http':
            keys = ['req_path', 'req_method', 'resp_status', 'resp_message']
            record_msg = cls.parse_msg(keys, data['record_msg'])
            name = f"{record_msg['req_method']} {record_msg['req_path']}"
            req_type = record_msg['req_method']
            req_status = str(record_msg['resp_status'])
            msg = cls.HttpMsg(req_method=record_msg['req_method'], req_path=record_msg['req_path'],
                              resp_status=record_msg['resp_status'], resp_message=record_msg['resp_message'])

        if data['protocol'] == 'redis':
            keys = ['_', 'req_command', 'req_payload', '_', '_', 'resp_command', 'resp_payload', '_']
            record_msg = cls.parse_msg(keys, data['record_msg'])
            name = record_msg['req_command']
            req_type = record_msg['req_command']
            if record_msg['resp_payload'][0] == '-':
                req_status = 'ERR'
            else:
                req_status = 'OK'
            msg = cls.RedisMsg(req_command=record_msg['req_command'], req_payload=record_msg['req_payload'],
                               resp_command=record_msg['resp_command'], resp_payload=record_msg['resp_payload'])

        return cls(name=name,
                   start_timestamp=data['req_timestamp'],
                   end_timestamp=data['resp_timestamp'],
                   children=list(), tgid=data['tgid'],
                   thread_id=data['thread_id'],
                   id=f"{data['tgid']}_{data['thread_id']}_{data['req_timestamp']}_{data['resp_timestamp']}",
                   client_ip=data['client_ip'],
                   server_ip=data['server_ip'],
                   server_port=data['server_port'],
                   client_port=data['client_port'],
                   protocol=data['protocol'],
                   value=int(data['resp_timestamp']) - int(data['req_timestamp']),
                   req_type=req_type,
                   req_status=req_status,
                   msg=msg,
                   raw_data=data
                   )


class SpanEncoder(JSONEncoder):
    def default(self, o):
        return o.__dict__


class SpanRelation:
    def __init__(self, raw_data1, raw_data2):
        # raw_data1 为tcp_client, raw_data2为tcp_server
        if raw_data1['l4_role'] == 'tcp_server' and raw_data1['l7_role'] == 'server':
            tmp = raw_data1
            raw_data1 = raw_data2
            raw_data2 = tmp
        self.parent_id = f"{raw_data1['tgid']}_{raw_data1['thread_id']}_{raw_data1['req_timestamp']}_{raw_data1['resp_timestamp']}"
        self.child_id = f"{raw_data2['tgid']}_{raw_data2['thread_id']}_{raw_data2['req_timestamp']}_{raw_data2['resp_timestamp']}"
        self.parent_data = raw_data1
        self.child_data = raw_data2


def check_timestamp_valid(span_relation):
    parent_span = span_relation.parent_data
    child_span = span_relation.child_data
    # 误差 3s
    return int(parent_span['req_timestamp']) <= int(child_span['req_timestamp']) + TIME_DEVIATION and \
           int(child_span['resp_timestamp']) <= int(parent_span['resp_timestamp']) + TIME_DEVIATION


def find_span_by_tgid_threadid_time(spans, parent_data):
    for span in spans.values():
        if span.tgid == parent_data['tgid'] and span.thread_id == parent_data['thread_id'] and \
                int(span.start_timestamp) - TIME_DEVIATION <= int(parent_data['req_timestamp']) <= int(span.end_timestamp) + TIME_DEVIATION and \
                int(span.start_timestamp) - TIME_DEVIATION <= int(parent_data['resp_timestamp']) <= int(span.end_timestamp) + TIME_DEVIATION:
            return span
    return None


def get_data(prometheus_server, start, end, step):
    prometheus_url = f'http://{prometheus_server}/api/v1/query_range?query=gala_gopher_l7_trace_value&start={start}&end={end}&step={step}'

    # 查询表达式
    query = 'gala_gopher_l7_trace_value'
    # 发送 GET 请求
    response = requests.get(prometheus_url, params={'query': query})

    # 解析 JSON 数据
    data = response.json()
    LOGGER.info(f"gala_gopher_l7_trace_value prometheus data is {data}")

    raw_metrics = list()
    for result in data['data']['result']:
        raw_metrics.append(result['metric'])

    spans = dict()
    for metric in raw_metrics:
        if metric['l4_role'] == 'tcp_server' and metric['l7_role'] == 'server':
            span = Span.parse_span(metric)
            spans[span.id] = span

    span_relations = list()
    for raw_span in combinations(raw_metrics, 2):
        if ((raw_span[0]['req_tcp_seq'] == raw_span[1]['req_tcp_seq'] and raw_span[0]['resp_tcp_seq'] == raw_span[1]['resp_tcp_seq'])
                or (raw_span[0]['record_digest'] == raw_span[1]['record_digest'] and raw_span[0]['l7_role'] != raw_span[1]['l7_role'])):
            span_relations.append(SpanRelation(raw_span[0], raw_span[1]))

    all_span = copy.copy(spans)
    LOGGER.info(f"gala_gopher_l7_trace_value SpanRelation is {json.dumps(list(all_span.values()), cls=SpanEncoder)}")
    for span_relation in span_relations:
        if check_timestamp_valid(span_relation):
            parent_span = find_span_by_tgid_threadid_time(spans, span_relation.parent_data)
            LOGGER.info(f"gala_gopher_l7_trace_value SpanRelation parent_span is {parent_span}")
            if parent_span is None:
                continue
            parent_span.children.append(spans[span_relation.child_id])
            if span_relation.child_id in all_span:
                all_span.pop(span_relation.child_id)

    for span_id in spans.keys():
        spans[span_id].children.sort(key=lambda x: x.start_timestamp)

    return json.dumps(list(all_span.values()), cls=SpanEncoder)


class DistributeTracingAPI(BaseResponse):
    @BaseResponse.handle(token=False)
    def get(self, **params):
        if not configuration.prometheus or not configuration.prometheus.host or not configuration.prometheus.port:
            return self.response(code=SUCCEED, data={})
        prometheus_server = f'{configuration.prometheus.host}:{configuration.prometheus.port}'
        now = time.time()
        step = 5
        # 默认展示10分钟内数据
        end = int(params.get("endTime", int(now) * 1000)) / 1000
        start = int(params.get("startTime", end * 1000 - TEN_MINUTE)) / 1000

        query_data = json.loads(get_data(prometheus_server, start, end, step))
        for data in query_data:
            data.pop('children')
            data.pop('raw_data')
        prom_query_params = {
            "start_timestamp": start,
            "end_timestamp": end,
            "step": step
        }
        res_data = {"query_params": prom_query_params, "request_list": query_data}
        return self.response(code=SUCCEED, data=res_data)

    @BaseResponse.handle(token=False)
    def post(self, **params):
        try:
            prom_query_params = params['query_params']
            trace_id = params['id']
            prometheus_server = f'{configuration.prometheus.host}:{configuration.prometheus.port}'
            query_data = json.loads(get_data(prometheus_server, prom_query_params['start_timestamp'], prom_query_params['end_timestamp'],
                                             prom_query_params['step']))
            for data in query_data:
                if data['id'] == trace_id:
                    return self.response(code=SUCCEED, data=data)
        except Exception as e:
            LOGGER.error(e)
        return self.response(code=SUCCEED, data="Non Found")
